"""
This module contains the code to parse the configuration file.
The configuration should be a key-value dictionary. The keys are templates/tasks names.

The key '_' is a magic key that is used to inherit the configuration from the parent node.

a task should be defined as a table:

__task__: <task_name>, # the entry point/function name of the task
kwarg1: <kwarg1_value>
kwarg2: <kwarg2_value>
...

Each value can be raw value or a reference to a file, or a reference to a task.

a reference to a file shoule be:
__file__: <file_id>

a reference to a task should be:
__task__: <task_hash>

a reference to a file generated by a task should be:
__task__: <task_name>
"""

import os
from typing import Dict, Any, List
from functools import cached_property
from pydantic import BaseModel
from loguru import logger

reserved_keys = ["__task__", "__root__"]


class DefaultValueNotProvided(Exception):
    pass


class _Node:
    """
    A node is a wrapper for a value that can be a dictionary or a value.
    It is used to handle the inheritance and special keys.
    """

    _parent: "_Node" = None
    _name: str = None
    _children: Dict[str, "_Node"] = None
    _value: Any = None

    def __init__(self, cfg, parent=None, name="ROOT"):
        self._parent = parent
        self._name = name
        self._raw = cfg
        self._valid = True
        if isinstance(cfg, dict):
            self._value = None
            self._children = {}
            for k, v in cfg.items():
                if k == "_":
                    continue
                self._children[k] = _Node(v, self, k if parent==None else None)
            inherit = cfg.get("_", None)
            if inherit is not None:
                if isinstance(inherit, str):
                    inherit = [inherit]
                try:
                    inherit_nodes = {}
                    for inherit_cfg in inherit:
                        inherit_value = _Node(
                            self._parent.search(inherit_cfg),
                            self,
                            f"__inherit__:{inherit_cfg}",
                        )
                        assert (
                            inherit_value._children is not None
                        ), "Inherit node must have children"
                        inherit_nodes.update(inherit_value._children)
                    for k, v in inherit_nodes.items():
                        if k not in self._children:
                            self._children[k] = v
                except KeyError:
                    self._valid = False
        else:
            self._value = cfg
            self._children = None

    @cached_property
    def valid(self):
        if not self._valid:
            return False
        if self._children is not None:
            for k, v in self._children.items():
                if not v.valid:
                    return False
        return True

    def dump(self) -> Dict[str, Any] | Any:
        if self._children is None:
            return self._value
        ret_value = {}
        for k, v in self._children.items():
            if not v.valid:
                continue
            ret_value[k] = v.dump()
        if self._name is not None:
            ret_value["__name__"] = self._name
        return ret_value

    def search(self, key):
        node = self
        ret_value = None
        while node is not None:
            if key in node._raw:
                ret_value = node._raw[key]
            node = node._parent
        if ret_value is None:
            raise KeyError(f"Key {key} not found")
        return ret_value

    def _print_kv(self, indent, k, v):
        prefix = "  " * indent
        if isinstance(v, _Node) and v._children is None:
            print(f"{prefix}{k}: {v._value}")
        else:
            print(f"{prefix}{k}:")
            v.print_tree(indent + 1)

    def print_tree(self, indent=0):
        for k, v in self._children.items():
            self._print_kv(indent, k, v)

    def __repr__(self):
        node = self
        name_chain = []
        while node is not None:
            if node._name is not None:
                name_chain.append(node._name)
            else:
                name_chain.append("")
            node = node._parent
        name_chain.reverse()
        name = "/".join(name_chain)
        return f"Node(name={name}, value={self._value})"


@logger.catch
def parse_cfg(cfg: Dict[str, Any], env: Dict[str, Any] | None = None) -> Any:
    """
    Parse the configuration file recursively, handling inheritance and special keys.

    Args:
        cfg: A dictionary containing configuration data. May contain special keys:
            '_': String key that references another key in the parent dict for inheritance
            '__task__': Indicates a task reference
            '__file__': Indicates a file reference

    Returns:
        Dict[str, Any]: Processed configuration with inheritance resolved and '_' keys removed
    """
    node = _Node(cfg)
    node.print_tree()
    cfg = node.dump()
    return cfg


if __name__ == "__main__":
    import tomli

    with open("aatask_mnist.toml", "rb") as f:
        cfg = tomli.load(f)
    parse_cfg(cfg)
